// Copyright 2011 Cooliris, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Compiling this file as .mm generates some C++ template errors and crashes the compiler
#import <opencv2/objdetect/objdetect.hpp>
#import <opencv2/core/core_c.h>

#import "ObjectDetector.h"

@interface Classifier : NSObject {
@private
  CvHaarClassifierCascade* _classifier;
  NSUInteger _minSize;
}
- (id) initWithPath:(NSString*)path mininumObjectSize:(NSUInteger)minSize;
@property(nonatomic, readonly) CvHaarClassifierCascade* classifier;
@property(nonatomic, readonly) NSUInteger mininumObjectSize;
@end

@implementation Classifier

@synthesize classifier=_classifier, mininumObjectSize=_minSize;

- (id) initWithPath:(NSString*)path mininumObjectSize:(NSUInteger)minSize {
  if ((self = [super init])) {
    _classifier = (CvHaarClassifierCascade*)cvLoad([path fileSystemRepresentation], NULL, NULL, NULL);
    if (_classifier == NULL) {
      [self release];
      return nil;
    }
    _minSize = minSize;
  }
  return self;
}

- (void) dealloc {
  if (_classifier) {
    cvReleaseHaarClassifierCascade(&_classifier);
  }
  
  [super dealloc];
}

@end

@implementation ObjectDetector

- (id) init {
  return [self initWithMaximumProcessingSize:0];
}

- (id) initWithMaximumProcessingSize:(NSUInteger)maxProcessingSize {
  if ((self = [super init])) {
    _maxProcessingSize = maxProcessingSize;
    _classifiers = [[NSMutableArray alloc] init];
  }
  return self;
}

- (void) dealloc {
  [_classifiers release];
  
  [super dealloc];
}

- (BOOL) loadClassifierFromXMLAtPath:(NSString*)classifierPath minimumObjectSize:(NSUInteger)minSize {
  Classifier* classifier = [[Classifier alloc] initWithPath:classifierPath mininumObjectSize:minSize];
  if (classifier) {
    [_classifiers addObject:classifier];
    [classifier release];
  }
  return classifier ? YES : NO;
}

- (NSArray*) _detectObjectsInCGImage:(CGImageRef)image {
  NSMutableArray* results = nil;
  if (_classifiers.count) {
    // Compute processing dimensions
    CGFloat imageWidth = CGImageGetWidth(image);
    CGFloat imageHeight = CGImageGetHeight(image);
    size_t contextWidth;
    size_t contextHeight;
    CGFloat scale;
    if ((_maxProcessingSize > 0) && ((imageWidth > _maxProcessingSize) || (imageHeight > _maxProcessingSize))) {
      CGSize size = CGSizeMake(_maxProcessingSize, _maxProcessingSize);
      if (imageWidth / size.width >= imageHeight / size.height) {
        scale = ((CGFloat)imageWidth / size.width);
        contextHeight = (CGFloat)imageHeight / scale;
        contextWidth = size.width;
      } else {
        scale = ((CGFloat)imageHeight / size.height);
        contextWidth = (CGFloat)imageWidth / scale;
        contextHeight = size.height;
      }
    } else {
      contextWidth = imageWidth;
      contextHeight = imageHeight;
      scale = 1.0;
    }
    
    // Create BGR IplImage from rendered CGImage
    IplImage* iplImage = cvCreateImage(cvSize(contextWidth, contextHeight), IPL_DEPTH_8U, 4);
    CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
    CGContextRef contextRef = CGBitmapContextCreate(iplImage->imageData, iplImage->width, iplImage->height, iplImage->depth,
                                                    iplImage->widthStep, colorSpace,
                                                    kCGImageAlphaPremultipliedLast | kCGBitmapByteOrderDefault);
    CGContextDrawImage(contextRef, CGRectMake(0, 0, contextWidth, contextHeight), image);
    CGContextRelease(contextRef);
    CGColorSpaceRelease(colorSpace);
    
    // Find objects
    results = [NSMutableArray array];
    for (Classifier* classifier in _classifiers) {
      CvMemStorage* storage = cvCreateMemStorage(0);
      CvSize minSize = cvSize(classifier.mininumObjectSize, classifier.mininumObjectSize);
      CvSeq* objects = cvHaarDetectObjects(iplImage, classifier.classifier, storage, 1.1, 3, CV_HAAR_DO_CANNY_PRUNING, minSize,
                                           cvSize(0, 0));
      NSMutableArray* array = [[NSMutableArray alloc] init];
      for (int i = 0; i < objects->total; ++i) {
        CvRect* objectRect = (CvRect*)cvGetSeqElem(objects, i);
        [array addObject:[NSValue valueWithCGRect:CGRectMake((CGFloat)objectRect->x * scale, (CGFloat)objectRect->y * scale,
                                                             (CGFloat)objectRect->width * scale, (CGFloat)objectRect->height * scale)]];
      }
      [results addObject:array];
      [array release];
      cvReleaseMemStorage(&storage);
    }
    
    // Clean up
    cvReleaseImage(&iplImage);
  }
  return results;
}

- (NSArray*) detectObjectsInImage:(UIImage*)image {
  return image ? [self _detectObjectsInCGImage:[image CGImage]] : nil;
}

@end
